package com.cet.designapplication.iterator;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.http.HttpResponse;
import cn.hutool.http.HttpUtil;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import com.alibaba.fastjson.JSON;

import javax.persistence.Id;
import java.util.*;

/**
 * @author 30952
 */
public class EsQueryCollection implements Iterable<Map<String, Object>>{

    private List<Map<String, Object>> dataList = new LinkedList<>();

    private final List<String> columnList = new LinkedList<>();

    private  String cursor;

    public EsQueryCollection(String querySql, Integer fetchSize) {
        JSONObject jsonObject = new JSONObject();
        jsonObject.set("query", querySql);
        jsonObject.set("fetchSize", fetchSize);
        // http://localhost:9200/_sql?format=json
        String result = HttpUtil.createPost("http://localhost:9200/_sql?format=json")
                .body(jsonObject.toString())
                .execute()
                .body();
        EsResponseData responseData = JSONUtil.toBean(result, EsResponseData.class);

        // 存储keys
        List<Map<String, String>> columns = responseData.getColumns();
        if (CollUtil.isEmpty(columns)) {
            throw new UnsupportedOperationException("没有相关数据！");
        }

        for (Map<String, String> column : columns) {
            columnList.add(column.get("name"));
        }

        // 存储初始数据
        List<List<Object>> rows = responseData.getRows();
        if (CollUtil.isEmpty(rows)) {
            return;
        }

        for (List<Object> row : rows) {
            if (columnList.size() != row.size()) {
                throw new RuntimeException("数据与列数不匹配！");
            }
            HashMap<String, Object> data = new HashMap<>();
            for (int i = 0; i < columnList.size(); i++) {
                data.put(columnList.get(i), row.get(i));
            }
            dataList.add(data);
        }

        // 存储游标
        this.cursor = responseData.getCursor();
    }

    @Override
    public Iterator<Map<String, Object>> iterator() {
        return new EsQueryIterator(cursor, columnList, dataList);
    }

    private static class EsQueryIterator implements Iterator<Map<String, Object>> {

        private String cursor;

        private final List<String> columnList;

        private final List<Map<String, Object>> dataList;

        private Integer index;

        public EsQueryIterator(String cursor, List<String> columnList, List<Map<String, Object>> dataList) {
            this.cursor = cursor;
            this.columnList = columnList;
            this.dataList = dataList;
            this.index = 0;
        }

        @Override
        public boolean hasNext() {
            if (index < dataList.size()) {
                return true;
            }
            if (cursor != null) {
                JSONObject jsonObject = new JSONObject();
                jsonObject.set("cursor", cursor);
                String result = HttpUtil.createPost("http://localhost:9200/_sql?format=json")
                        .body(jsonObject.toString())
                        .execute()
                        .body();
                EsResponseData responseData = JSONUtil.toBean(result, EsResponseData.class);
                if (CollUtil.isEmpty(responseData.getRows())) {
                    return false;
                }
                for (List<Object> row : responseData.getRows()) {
                    if (columnList.size() != row.size()) {
                        throw new RuntimeException("数据与列数不匹配！");
                    }
                    HashMap<String, Object> data = new HashMap<>();
                    for (int i = 0; i < columnList.size(); i++) {
                        data.put(columnList.get(i), row.get(i));
                    }
                    dataList.add(data);
                }

                // 存储游标
                this.cursor = responseData.getCursor();
                return true;
            }
            return false;
        }

        @Override
        public Map<String, Object> next() {
            return dataList.get(index++);
        }
    }
}
